#include "ldecode.hpp"

//Notes:
//vSize = N
//vNodes: 1, 2, ..., N.
//cSize = M
//check Nodes: 1, 2, ..., M


namespace LDPCDecoder{
  // compute the index of qmatrix
  inline long qindex(long& i, long& j, long& m){
    return (i * (m+1) + j);
  }

  // compute the index of rmatrix
  inline long rindex(long& i, long& j, long& n){
    return (i * (n+1) + j);
  }

  // initialization of q_ij(0), q_ij(1)s are implied by 1 - q_ij(0).
  void init_variable_node(double* qmat, const double* P, const mat_GF2& H){
    long n = H.NumCols();
    long m = H.NumRows();
    for(long v = 1; v <= n; v++){
      const double qij = 1 - P[v];
      for(long c = 1; c <= m; c++){
	if (IsOne(H(c, v))) // then c is a neighbor of v
	  qmat[qindex(v, c, m)] = qij;
      }
    }
  }

  // compute response from check node c_j. i.e., r_ji(0)
  double comp_check_node_response(const double* qmat, 
				  long& i, long& j, 
				  const mat_GF2& H, 
				  const double& multresult){

    long n = H.NumCols();
    long m = H.NumRows();
    const double a = 2 * qmat[qindex(i, j, m)] - 1;
    if (a > 0)
      return (0.5 * (multresult/a) + 0.5);
    double result = 1;

    for(long v = 1; v <= n; v++){
      if ((v != i) && (IsOne(H(j, v)))) {
	result = result * (2 * qmat[qindex(v, j, m)] - 1);
      }
    }
    return 0.5 * (result + 1);
  }

  // obtain the neighbors of the variable node v.
  vector<long>* get_neighbors_v(const mat_GF2& H, long& v) {
    vector<long>* result = new vector<long>;
    long m = H.NumRows();
    for(long c = 1; c <= m; c++){
      if (IsOne(H(c, v))){
	result -> push_back(c);
      }
    }
    return result;
  }

  // obtain the neighbors of the check node c.
  vector<long>* get_neighbors_c(const mat_GF2& H, long& c) {
    vector<long>* result = new vector<long>;
    long n = H.NumCols();
    for(long v = 1; v <= n; v++){
      if (IsOne(H(c, v))){
	result -> push_back(v);
      }
    }
    return result;
  }

  // check whether C is a codeword by verifying H * C = 0
  inline bool is_code_word(const vec_GF2& C, const mat_GF2& H){
    GF2 x;
    long m = H.NumRows();
    for(long i = 1; i <= m; i++){
      InnerProduct(x, H(i), C);
      if (IsOne(x)) 
	return false;
    }
    return true;
  }

  inline double multiply_response1(const double* rmat, long& i, long& j, long& n,
				   vector<long>* neighbors, const double& pi){
    double result = 1.0;
    for(long i = 0; i < (long) neighbors -> size(); i++){
      long c = (*neighbors)[i];
      if (c != j)
	result *= rmat[rindex(c, i, n)];
    }
    return (1-pi) * result;
  }

  inline double multiply_response2(const double* rmat, long& i, long& j, long& n,
				   vector<long>* neighbors, const double& pi){
    double result = 1.0;
    for(long i = 0; i < (long) neighbors -> size(); i++){
      long c = (*neighbors)[i];
      if (c != j)
	result *= 1 - rmat[rindex(c, i, n)];
    }
    return pi * result;
  }

  // compute updates to qmat(i, j) for decoding step 4.
  double compute_variable_node_response(const double* rmat, 
					long& i, long& j, long& n,
					vector<long>* neighbors, 
					const double& pi, const double& Q0i, 
					const double& Q1i){
    if ((pi == 1.0) || (pi == 0.0)) 
      return (1.0 - pi);
    double coef0 = Q0i;
    double coef1 = Q1i;
    double rmatji = rmat[rindex(j, i, n)];
    if (rmatji == 0.0) {
      coef0 = multiply_response1(rmat, i, j, n, neighbors, pi);
    }
    else if (rmatji == 1.0) {
      coef1 = multiply_response2(rmat, i, j, n, neighbors, pi);
    }
    else{
      coef0 = coef0/rmatji;
      coef1 = coef1/(1 - rmatji);
    }
    return (1/(coef0 + coef1)) * coef0;
  }

  // free the arrays used in the decoding.
  inline void free_arrays(double* qmat, double* rmat, double* Q0, double* Q1){
    delete [] qmat;
    delete [] rmat; 
    delete [] Q0;
    delete [] Q1;
  }


  inline void compute_multiplies(double& multresult, const double* qmat, const mat_GF2& H, long& c){
    multresult = 1;
    double qvc;
    long n = H.NumCols();
    long m = H.NumRows();
    for(long v = 1; v <= n; v++){
      if (IsOne(H(c, v))){
	qvc = qmat[qindex(v, c, m)];
	if (qvc != 1)
	  multresult = multresult * (2 * qvc - 1);
      }
    }
  }


  inline void compute_rji(double* rmat, const double* qmat, const mat_GF2& H){
    double multresult;
    long m = H.NumRows();
    long n = H.NumCols();
    for(long c = 1; c <= m; c++){
      compute_multiplies(multresult, qmat, H, c);
      for(long v = 1; v <= n; v++){
	if (IsOne(H(c, v))){      
	  rmat[rindex(c, v, n)] = comp_check_node_response(qmat, v, c, H, multresult);
	}
      }
    }
  }

  // This function implements the LDPC decoder using the sum product
  // algorithm in probability domain.
  // C: to store the corrected codeword
  // H: parity-check matrix .
  // iter: to store the number of iterations used in decoding
  // P: a list of initial error probabilities.
  bool decodeSPA(vec_GF2& C, const mat_GF2& H, int& iter, const double* P){
    iter = 0;
    long n = H.NumCols();
    long m = H.NumRows();

    C.SetLength(n);

    double* qmat = new double[(n+1) * (m+1)];
    double* rmat = new double[(m+1) * (n+1)];
    double* Q0 = new double[n+1];
    double* Q1 = new double[n+1];
    double threshold;

    // 1. initialization of q_ij(0)
    init_variable_node(qmat, P, H);

    while(iter < MAX_ITER){
#ifdef VERBOSE
      cout << "Iteration: " << iter << endl;
#endif
      // 2. compute r_ji(0), where check node c_j and variable node v_i
      // are neighbors.
      compute_rji(rmat, qmat, H);

      // 3. compute the estimation of codewords.
      for(long v = 1; v <= n; v++) {
	vector<long>* neighbors = get_neighbors_v(H, v);
	double prod0 = 1.0 - P[v];
	double prod1 = 1.0 - prod0;
	if (prod0 == 0.0) {
	  for(long i = 0; i < (long) neighbors -> size(); i++){
	    long c = (*neighbors)[i];
	    prod1 = prod1 * (1 - rmat[rindex(c, v, n)]);
	  }
	}
	else if (prod1 == 0.0){
	  for(long i = 0; i < (long) neighbors -> size(); i++){
	    long c = (*neighbors)[i];
	    prod0 = prod0 * rmat[rindex(c, v, n)];
	  }	
	}
	else {
	  for(long i = 0; i < (long) neighbors -> size(); i++){
	    long c = (*neighbors)[i];
	    prod0 = prod0 * rmat[rindex(c, v, n)];
	    prod1 = prod1 * (1 - rmat[rindex(c, v, n)]);
	  }
	}
	delete neighbors;
	Q0[v] = prod0;
	Q1[v] = prod1;
	threshold = prod1/(prod0 + prod1);
	if (threshold > 0.5)
	  C(v) = to_GF2(1);
	else if (threshold < 0.5)
	  C(v) = to_GF2(0);
	else{
	  cout << threshold << endl;
	  cout << "Overflow may happen!" << endl;
	  delete neighbors;
	  free_arrays(qmat, rmat, Q0, Q1);
	  exit(EXIT_FAILURE);
	}
      }

      // if decoding is successful, terminate
      if (is_code_word(C, H)){
#ifdef VERBOSE
	cout << "Success!" << endl;
#endif
	free_arrays(qmat, rmat, Q0, Q1);
	return true;
      }
        
      // 4. variable nodes update the response messages q_ij to check
      // nodes via equation solving.
      for(long v = 1; v <= n; v++){
	vector<long>* neighbors = get_neighbors_v(H, v);
	for(long c = 1; c <= m; c++){
	  qmat[qindex(v, c, m)] = compute_variable_node_response(rmat, v, c, n, neighbors, P[v], Q0[v], Q1[v]);
	}
	delete neighbors;
      }
      iter++;
    }

#ifdef VERBOSE
    cout << "Warning: Nonconverge!" << endl;
#endif
    free_arrays(qmat, rmat, Q0, Q1);
    return false;
  }


}
